import numpy as np
from model.utils import acc_dist
import heapq


def get_dict_fmt(order, topk_drivers):
    # 生成字典格式
    list = []
    for driver in topk_drivers:
        dict_fmt = {}
        dict_fmt["order_id"] = int(order[0])
        dict_fmt["driver_id"] = int(driver[0])
        order_lng, order_lat = float(order[3]), float(order[4])
        driver_lng, driver_lat = float(driver[2]), float(driver[3])
        dist = acc_dist(order_lng, order_lat, driver_lng, driver_lat)
        dict_fmt["order_driver_distance"] = dist
        order_start_location = [float(order[3]), float(order[4])]
        dict_fmt["order_start_location"] = order_start_location
        order_finish_location = [float(order[5]), float(order[6])]
        dict_fmt["order_finish_location"] = order_finish_location
        driver_location = [float(driver[2]), float(driver[3])]
        dict_fmt["driver_location"] = driver_location
        dict_fmt["timestamp"] = int(order[1])
        dict_fmt["order_finish_timestamp"] = int(order[2])
        dict_fmt["day_of_week"] = int(2)
        dict_fmt["reward_units"] = float(order[7])
        dict_fmt["pick_up_eta"] = float(0)

        list.append(dict_fmt)

    return list

def get_topk(order_skill, order, drivers):
    """
    为一个order从同一时间片的drivers中选取top-11个最近的工人
    :param order_skill: 从外部指定，因为有的order技能为3，则同时需要技能1和技能2
    :param order: 订单
    :param drivers: 同一时间片的司机list
    :return: top-11个技能匹配且离order最近的drivers
    """
    # 选取有效工人
    valid_drivers = []
    for driver in drivers:
        # 技能匹配
        if driver[4] == order_skill:
            valid_drivers.append(driver)

    # Top-K算法：使用最小堆
    valid_drivers_top11 = []
    heap = []
    # 订单的经、纬度
    order_lng, order_lat = float(order[3]), float(order[4])

    for i in range(len(valid_drivers)):
        driver_lng, driver_lat = float(valid_drivers[i][2]), float(valid_drivers[i][3])
        dist = acc_dist(order_lng, order_lat, driver_lng, driver_lat)
        heapq.heappush(heap, (dist, i))

    valid_drivers_top11_indexs = []
    i = 0
    while len(heap) > 0 and i < 11:
        _, index = heapq.heappop(heap)
        valid_drivers_top11_indexs.append(index)
        i += 1

    # 堆中保存了Top11个Driver再drivers中的index
    # 再把对应的信息保存并返回
    for index in valid_drivers_top11_indexs:
        valid_drivers_top11.append(valid_drivers[index])

    return valid_drivers_top11
